import copy
import os
import pickle
from tensorflow.keras.utils import to_categorical
import numpy as np
import matplotlib.pyplot as plt

# MAX full training time
MAX_FTT = 5521.803059895833
# MAX trainable parameters
MAX_TP = 49979274

def delete_margin(matrix):
    return matrix[:-1, 1:]


# return: 1. X: a linear array: flattened adjacent matrix + integer operations
#         2. y: accuracy
def get_toy_data(important_metrics, select_upper_tri=False, max_creation=-1,
                 integers2one_hot=True, additional_metrics=False, normalization=True):
    X = []
    y = []
    more_metrics_num = []
    for index in important_metrics:
        fixed_metrics = important_metrics[index]['fixed_metrics']
        adjacent_matrix = fixed_metrics['module_adjacency']
        module_integers = fixed_metrics['module_integers']
        trainable_parameters = fixed_metrics['trainable_parameters']
        final_training_time = important_metrics[index]['final_training_time']
        accuracy = important_metrics[index]['final_test_accuracy']

        adjacent_matrix = delete_margin(adjacent_matrix)
        flattened_adjacent = adjacent_matrix.flatten()
        input_metrics = []
        input_metrics.extend(flattened_adjacent)
        if integers2one_hot:
            module_integers = to_categorical(module_integers, 4, dtype='int8')
            module_integers = module_integers.flatten()
        input_metrics.extend(module_integers)
        if additional_metrics:
            if normalization:
                trainable_parameters = trainable_parameters / MAX_TP
                final_training_time = final_training_time / MAX_FTT
            input_metrics.extend([trainable_parameters])
            input_metrics.extend([final_training_time])
        X.append(input_metrics)
        y.append(accuracy)

    assert len(X) == len(y)
    print('Input {:} metrics, obtain {:} metrics'.format(len(important_metrics), len(X)))
    return X, y, more_metrics_num


def get_upper_triangular_data(important_metrics, integers2one_hot=True, double_upper=True, additional_metrics=True,
                              normalization=True):
    # upper triangular and additional metrics (including trainable parameters and final training time)
    # double_upper denotes flatting the upper triangular matrix into a one-dimensional vector based the axis = 0 and 1
    X = []
    y = []
    for index in important_metrics:
        fixed_metrics = important_metrics[index]['fixed_metrics']
        adjacent_matrix = fixed_metrics['module_adjacency']
        module_integers = fixed_metrics['module_integers']
        if integers2one_hot:
            module_integers = to_categorical(module_integers, 4, dtype='int8')
            module_integers = module_integers.flatten()

        trainable_parameters = fixed_metrics['trainable_parameters']
        final_training_time = important_metrics[index]['final_training_time']
        adjacent_matrix = get_data.delete_margin(adjacent_matrix)
        array_adjacent_matrix = np.array(adjacent_matrix)
        flattened_adjacent = []
        matrix_size = len(adjacent_matrix)
        # get upper triangular data in matrix and flat it
        # Noting that this doesn't contains the elements as main diagonal
        for i in range(matrix_size):
            flattened_adjacent.extend(adjacent_matrix[i][i:])

        if double_upper:
            # add information from column
            for i in range(matrix_size):
                extend_array = array_adjacent_matrix[:i + 1, i]
                flattened_adjacent.extend(extend_array.tolist())
        input_metrics = []
        input_metrics.extend(flattened_adjacent)
        input_metrics.extend(module_integers)
        if additional_metrics:
            if normalization:
                trainable_parameters = trainable_parameters / MAX_TP
                final_training_time = final_training_time / MAX_FTT
            input_metrics.extend([trainable_parameters])
            input_metrics.extend([final_training_time])
        accuracy = important_metrics[index]['final_test_accuracy']
        X.append(input_metrics)
        y.append(accuracy)

    return X, y


def get_toy_metrics(num, type='train', train_num=2000):
    index = get_data.get_data_index_from_101(num, type=type, train_num=train_num)
    metrics = get_data.get_corresponding_metrics_by_index(index, type=type)   # 获取结构，精度训练时间信息
    metrics = get_data.padding_zero_in_matrix(metrics)   # 矩阵和操作补长定长
    metrics = get_data.operations2integers(metrics)     # 操作转数字
    return metrics
